#include "corner_pool_plugin.h"
#include "cuda_fp16.h"


template<typename T>
__global__ void CornerPoolingWidthKernel(const T *pData,
                                         T *pOut,
                                         int64_t total_idx,
                                         int64_t w,
                                         int64_t cur_start_dim, int64_t next_start_dim, int64_t dim_n)
{
    for (int64_t w_idx = blockIdx.x * blockDim.x + threadIdx.x; w_idx < total_idx; w_idx += blockDim.x * gridDim.x)
    {
        T *pMaxVal = pOut + w_idx * w;
        const T *pCurVal = pData + w_idx * w;
        const T *pNextVal = pData + w_idx * w;
        for (int64_t i = 0; i < dim_n; i++)
        {
            T curVal = *(pCurVal + cur_start_dim + i);
            T nextVal = *(pNextVal + next_start_dim + i);
            if (curVal > nextVal)
            {
                *(pMaxVal + cur_start_dim + i) = curVal;
            }
            else
            {
                *(pMaxVal + cur_start_dim + i) = nextVal;
            }
        }
    }
}

template<typename T>
__global__ void CornerPoolingHeightKernel(const T *pData,
                                          T *pOut,
                                          int64_t n,
                                          int64_t c,
                                          int64_t h,
                                          int64_t w,
                                          int64_t cur_start_dim, int64_t next_start_dim, int64_t dim_n)
{
    int64_t total_idx = n * c;
    for (int64_t nc = blockIdx.x * blockDim.x + threadIdx.x; nc < total_idx; nc += blockDim.x * gridDim.x)
    {
        T *pMaxVal = pOut + nc * h * w + cur_start_dim * w;
        const T *pCurVal = pData + nc * h * w + cur_start_dim * w;
        const T *pNextVal = pData + nc * h * w + next_start_dim * w;
        for (int h_idx = 0; h_idx < dim_n; h_idx++)
        {
            for (int64_t i = 0; i < w; i++)
            {
                T curVal = *(pCurVal + h_idx * w + i);
                T nextVal = *(pNextVal + h_idx * w + i);
                if (curVal > nextVal)
                {
                    *(pMaxVal + h_idx * w + i) = curVal;
                }
                else
                {
                    *(pMaxVal + h_idx * w + i) = nextVal;
                }
            }
        }
    }
}

namespace
{
    const int KCudaNumThreads = 512;
}

template<typename T>
void CornerPoolingWidth(const T *pInput,
                        T *pOut,
                        int blockSize,
                        int64_t total_idx,
                        int64_t width,
                        int64_t start_dim0,
                        int64_t start_dim1,
                        int64_t dim_len,
                        cudaStream_t stream)
{
    CornerPoolingWidthKernel<T> <<<blockSize, KCudaNumThreads, 0, stream>>>(
            pInput, pOut, total_idx, width, start_dim0, start_dim1, dim_len);
}

template<typename T>
void CornerPoolingHeight(const T *pInput,
                         T *pOut,
                         int blockSize,
                         int64_t batch,
                         int64_t channel,
                         int64_t height,
                         int64_t width,
                         int64_t start_dim0,
                         int64_t start_dim1,
                         int64_t dim_len,
                         cudaStream_t stream)
{
    CornerPoolingHeightKernel<T> <<<blockSize, KCudaNumThreads, 0, stream>>>(
            pInput, pOut,
            batch, channel, height, width, start_dim0, start_dim1, dim_len);
}

template<typename T>
void LeftPoolingImpl(const void *const pInput,
                     void *pOut,
                     const int64_t batch,
                     const int64_t channel,
                     const int64_t height,
                     const int64_t width,
                     const nvinfer1::DataType type,
                     cudaStream_t stream)
{
    int block_size = (1 * channel * height + KCudaNumThreads - 1) / KCudaNumThreads;
    cudaMemcpyAsync(pOut, pInput, batch * channel * height * width * sizeof(T), cudaMemcpyDeviceToDevice,
                    stream);
    for (int64_t n = 0; n < batch; n++)
    {
        for (int64_t ind = 1; ind < width; ind <<= 1)
        {
            CornerPoolingWidth<T>(static_cast<const T *>(pOut),
                                      static_cast<T *>(pOut),
                                      block_size, 1 * channel * height, width, 0, ind, width - ind,
                                      stream);
        }
        pOut = (uint8_t *) pOut + channel * height * width * sizeof(T);
    }
}

void LeftPooling(const void *const pInput,
                 void *pOut,
                 const int64_t batch,
                 const int64_t channel,
                 const int64_t height,
                 const int64_t width,
                 const nvinfer1::DataType type,
                 cudaStream_t stream)
{
    if (type == nvinfer1::DataType::kFLOAT)
    {
        LeftPoolingImpl<float>(pInput, pOut, batch, channel, height, width, type, stream);
    }
    else if (type == nvinfer1::DataType::kINT8)
    {
        LeftPoolingImpl<int8_t>(pInput, pOut, batch, channel, height, width, type, stream);
    }
}

template<typename T>
void RightPoolingImpl(const void *const pInput,
                      void *pOut,
                      const int64_t batch,
                      const int64_t channel,
                      const int64_t height,
                      const int64_t width,
                      const nvinfer1::DataType type,
                      cudaStream_t stream)
{
    int block_size = (1 * channel * height + KCudaNumThreads - 1) / KCudaNumThreads;

    cudaMemcpyAsync(pOut, pInput, batch * channel * height * width * sizeof(T), cudaMemcpyDeviceToDevice,
                    stream);
    for (int64_t n = 0; n < batch; n++)
    {
        for (int64_t ind = 1; ind < width; ind <<= 1)
        {
            CornerPoolingWidth<T>(static_cast<const T *>(pOut),
                                  static_cast<T *>(pOut),
                                  block_size, 1 * channel * height, width, ind, 0, width - ind,
                                  stream);
        }
        pOut = (uint8_t *) pOut + channel * height * width * sizeof(T);
    }

}

void RightPooling(const void *const pInput,
                  void *pOut,
                  const int64_t batch,
                  const int64_t channel,
                  const int64_t height,
                  const int64_t width,
                  const nvinfer1::DataType type,
                  cudaStream_t stream)
{
    if (type == nvinfer1::DataType::kFLOAT)
    {
        RightPoolingImpl<float>(pInput, pOut, batch, channel, height, width, type, stream);
    }
    else if (type == nvinfer1::DataType::kINT8)
    {
        RightPoolingImpl<int8_t>(pInput, pOut, batch, channel, height, width, type, stream);
    }
}

template<typename T>
void TopPoolingImpl(const void *const pInput,
                    void *pOut,
                    const int64_t batch,
                    const int64_t channel,
                    const int64_t height,
                    const int64_t width,
                    const nvinfer1::DataType type,
                    cudaStream_t stream)
{
    int block_size = (1 * channel + KCudaNumThreads - 1) / KCudaNumThreads;
    cudaMemcpyAsync(pOut, pInput, batch * channel * height * width * sizeof(T), cudaMemcpyDeviceToDevice,
                    stream);
    for (int64_t n = 0; n < batch; n++)
    {
        for (int64_t ind = 1; ind < height; ind <<= 1)
        {
            CornerPoolingHeight<T>(static_cast<const T *>(pOut),
                                   static_cast<T *>(pOut),
                                   block_size, 1, channel, height, width, 0, ind, height - ind,
                                   stream);
        }
        pOut = (uint8_t *) pOut + channel * height * width * sizeof(T);
    }
}

void TopPooling(const void *const pInput,
                void *pOut,
                const int64_t batch,
                const int64_t channel,
                const int64_t height,
                const int64_t width,
                const nvinfer1::DataType type,
                cudaStream_t stream)
{
    if (type == nvinfer1::DataType::kFLOAT)
    {
        TopPoolingImpl<float>(pInput, pOut, batch, channel, height, width, type, stream);
    }
    else if (type == nvinfer1::DataType::kINT8)
    {
        TopPoolingImpl<int8_t>(pInput, pOut, batch, channel, height, width, type, stream);
    }
}

template<typename T>
void BottomPoolingImpl(const void *const pInput,
                       void *pOut,
                       const int64_t batch,
                       const int64_t channel,
                       const int64_t height,
                       const int64_t width,
                       const nvinfer1::DataType type,
                       cudaStream_t stream)
{
    int block_size = (1 * channel + KCudaNumThreads - 1) / KCudaNumThreads;
    cudaMemcpyAsync(pOut, pInput, batch * channel * height * width * sizeof(T), cudaMemcpyDeviceToDevice,
                    stream);
    for (int64_t n = 0; n < batch; n++)
    {
        for (int64_t ind = 1; ind < height; ind <<= 1)
        {
            CornerPoolingHeight<T>(static_cast<const T *>(pOut),
                                   static_cast<T *>(pOut),
                                   block_size, 1, channel, height, width, ind, 0, height - ind,
                                   stream);
        }
        pOut = (uint8_t *) pOut + channel * height * width * sizeof(T);
    }
}

void BottomPooling(const void *const pInput,
                   void *pOut,
                   const int64_t batch,
                   const int64_t channel,
                   const int64_t height,
                   const int64_t width,
                   const nvinfer1::DataType type,
                   cudaStream_t stream)
{
    if (type == nvinfer1::DataType::kFLOAT)
    {
        BottomPoolingImpl<float>(pInput, pOut, batch, channel, height, width, type, stream);
    }
    else if (type == nvinfer1::DataType::kINT8)
    {
        BottomPoolingImpl<int8_t>(pInput, pOut, batch, channel, height, width, type, stream);
    }
}

nvinfer1::PluginFieldCollection CornerPoolPluginCreator::plugin_field_collection = nvinfer1::PluginFieldCollection();

//REGISTER_TENSORRT_PLUGIN(CornerPoolPluginCreator);
